import cv2 as cv
import numpy as np
import rospy
from prometheus_msgs.msg import DetectionInfo
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image


detect = DetectionInfo()
bridge = CvBridge()
cv_img = None

def img_cb(msg):
    global cv_img
    cv_img = bridge.imgmsg_to_cv2(msg, "bgr8")

rospy.init_node("detect_red_point")
rate = rospy.Rate(20)
image_sub = rospy.Subscriber("/iris1/camera/image_raw", Image, img_cb, queue_size=5)
detect_pub = rospy.Publisher("/prometheus/object_detection/firepoint", DetectionInfo, queue_size=10)


while not rospy.is_shutdown():
    if cv_img is None:
        rospy.loginfo_once("waiting for img!")
        continue
    rospy.loginfo_once("Received image! start detection!")
    frame = cv_img.copy()
    shape = frame.shape
    frame = cv.medianBlur(frame, 5)

    hsv = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
    low_range = np.array([0, 100, 100])
    high_range = np.array([20, 255, 255])
    thresh = cv.inRange(hsv, lowerb=low_range, upperb=high_range)

    contours, hierarchy = cv.findContours(thresh, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    if len(contours) > 0:
        areas = [cv.contourArea(contour) for contour in contours]
        max_index = areas.index(max(areas))

        cv.drawContours(frame, contours, max_index, (0, 255, 255), 3)
        (x,y),radius = cv.minEnclosingCircle(contours[max_index])

        center = (int(x),int(y))
        radius = int(radius)
        cv.circle(frame,center,radius,(0,255,0),2)
        # cv.putText(frame, "center:[{0},{1}]".format(*center), center, 1, 1, (0, 0, 0), 1)

        detect.object_name = "red_point"
        detect.detected = True
        detect.position[0] = (center[0] - (shape[1] / 2)) / (shape[1] / 2)
        detect.position[1] = (center[1] - (shape[0] / 2)) / (shape[0] / 2)
        cv.putText(frame, "center:[{1},{0}]".format(detect.position[0], detect.position[1]), center, 1, 1, (0, 0, 0), 1)
    else:
        # rospy.loginfo("not detected!!!")
        detect.detected = False

    detect_pub.publish(detect)
    # rate.sleep()
    cv.imshow('frame', frame)
    cv.waitKey(1)

